package com.bigdata.spark.core.rdd.part

import org.apache.spark.rdd.RDD
import org.apache.spark.{Partitioner, SparkConf, SparkContext}

object RDD_Part {

  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().setMaster("local").setAppName("Partition")
    val sc = new SparkContext(sparkConf)

    val rdd = sc.makeRDD(List(
      ("zhenzhen", "aaaaaaa"),
      ("xianxian", "aaaaaaa"),
      ("hanhan", "aaaaaaa"),
      ("amaoamao", "aaaaaaa")
    ), 3)

    val partRDD: RDD[(String, String)] = rdd.partitionBy(new MyPartitioner)

    partRDD.saveAsTextFile("output")

    sc.stop()
  }

  /**
   * 自定义分区器
   * 1. 继承Partitioner
   * 2. 重写方法
   */
  class MyPartitioner extends Partitioner {
    // 分区数量
    override def numPartitions: Int = 3

    // 根据数据的key值返回数据的分区索引（从0开始）
    override def getPartition(key: Any): Int = {

      key match {
        case "xianxian" => 0
        case "amaoamao" => 1
        case _ => 2
      }

    }
  }
}
